#!/usr/bin/env python3

import rospy
from sensor_msgs.msg import JointState
import pygame
from pygame.locals import *
import sys

class KeyboardController:
    def __init__(self):
        rospy.init_node('key_controller')
        self.pub = rospy.Publisher('/joint_states', JointState, queue_size=10)
        
        # 关节名称及初始位置
        self.joint_names = [
            'joint_p1', 'joint_r1', 'joint_r2',
            'joint_r3', 'joint_r4', 'joint_r5', 'joint_r6'
        ]
        self.current_positions = {name: 0.0 for name in self.joint_names}
        
        # 按键配置（正转/反转）
        self.key_bindings = {
            K_1: ('joint_p1', +1),
            K_q: ('joint_p1', -1),
            K_2: ('joint_r1', +1),
            K_w: ('joint_r1', -1),
            K_3: ('joint_r2', +1),
            K_e: ('joint_r2', -1),
            K_4: ('joint_r3', +1),
            K_r: ('joint_r3', -1),
            K_5: ('joint_r4', +1),
            K_t: ('joint_r4', -1),
            K_6: ('joint_r5', +1),
            K_y: ('joint_r5', -1),
            K_7: ('joint_r6', +1),
            K_u: ('joint_r6', -1),
        }
        
        # 初始化PyGame
        pygame.init()
        self.screen = pygame.display.set_mode((200, 200))
        pygame.display.set_caption("ROS Arm Controller")
        self.clock = pygame.time.Clock()
        
        # 控制参数
        self.delta = 0.05  # 单次按键变化量（弧度）
        self.rate = 30     # 控制频率Hz

    def run(self):
        while not rospy.is_shutdown():
            # 处理系统事件
            for event in pygame.event.get():
                if event.type == QUIT:
                    pygame.quit()
                    sys.exit()
                elif event.type == KEYDOWN and event.key == K_ESCAPE:
                    pygame.quit()
                    sys.exit()

            # 获取键盘状态并更新关节位置
            pressed = pygame.key.get_pressed()
            for key in self.key_bindings:
                if pressed[key]:
                    joint, direction = self.key_bindings[key]
                    self.current_positions[joint] += self.delta * direction

            # 发布关节状态
            msg = JointState()
            msg.header.stamp = rospy.Time.now()
            msg.name = self.joint_names
            msg.position = [self.current_positions[name] for name in self.joint_names]
            self.pub.publish(msg)

            # 控制循环频率
            self.clock.tick(self.rate)

if __name__ == '__main__':
    try:
        controller = KeyboardController()
        controller.run()
    except rospy.ROSInterruptException:
        pass

# 写一个 key_controller.py 来充当 joint_state_publisher . 启动时打开一个窗口.当按下按键时发布 joint 控制机器人机械臂转动.关节名称和按键对应关系如下.
# | 关节名 | 正转 | 反转 |
# | --- | --- | --- |
# | joint_p1 | 1 | q |
# | joint_r1 | 2 | w |
# | joint_r2 | 3 | e |
# | joint_r3 | 4 | r |
# | joint_r4 | 5 | t |
# | joint_r5 | 6 | y |
# | joint_r6 | 7 | u |

